#pragma once


#define INT_DIV_CEIL(n, d) ((n - 1) / d) + 1

//This function expects to be called with blocksize = (16, 16)
__global__ void copy_k(Matrix input, Matrix output) {
  __shared__ float scratch[17*16];
  float* source_data = input.sub.data;
  float* dest_data = output.sub.data;
  bool input_col = (input.stride == 1);
  bool output_col = (output.stride == 1);
  if (input_col) {
    if (output_col) {
      // Both input and output in column major format
      int row = blockIdx.x * blockDim.x + threadIdx.x;
      int col = blockIdx.y * blockDim.y + threadIdx.y;
      if ((row < input.length) &&
          (col < input.sub.length)) {
        float source = *(source_data + input.sub.stride * col + row);
        *(dest_data + output.sub.stride * col + row) = source;
      }
    } else {
      // Input in Column major, output in row major
      int row = blockIdx.x * blockDim.x + threadIdx.x;
      int col = blockIdx.y * blockDim.y + threadIdx.y;
      if ((row < input.length) &&
          (col < input.sub.length)) {
        float source = *(source_data + input.sub.stride * col + row);
        scratch[threadIdx.y * 17 + threadIdx.x] = source;
      }
      int output_row = blockIdx.x * blockDim.x + threadIdx.y;
      int output_col = blockIdx.y * blockDim.y + threadIdx.x;
      int output_index = output_row * output.stride + output_col;
      __syncthreads();
      if ((output_row < output.length) &&
          (output_col < output.sub.length)) {
        *(dest_data + output_index) = scratch[threadIdx.x * 17 + threadIdx.y];
      }
    }
  } else {
    if (output_col) {
      // Input in Row major, output in column major
      int row = blockIdx.y * blockDim.y + threadIdx.y;
      int col = blockIdx.x * blockDim.x + threadIdx.x;
      if ((row < input.length) &&
          (col < input.sub.length)) {
        float source = *(source_data + input.stride * row + col);
        scratch[threadIdx.y * 17 + threadIdx.x] = source;
      }
      int output_row = blockIdx.y * blockDim.y + threadIdx.x;
      int output_col = blockIdx.x * blockDim.x + threadIdx.y;
      int output_index = output_col * output.sub.stride + output_row;
      __syncthreads();
      if ((output_row < output.length) &&
          (output_col < output.sub.length)) {
        *(dest_data + output_index) = scratch[threadIdx.x * 17 + threadIdx.y];
      }
    } else {
      // Input in Row major, output in Row major
      int col = blockIdx.x * blockDim.x + threadIdx.x;
      int row = blockIdx.y * blockDim.y + threadIdx.y;
      if ((row < input.length) &&
          (col < input.sub.length)) {
        float source = *(source_data + input.stride * row + col);
        *(dest_data + output.stride * row + col) = source;
      }
    }
  }                           
}

/**
 * Simple routine to copy matrix data on the device.  Will perform a transpose
 * if it detects a switch from row major to column major storage in the strides
 * of the arguments.  It expects to be called with the out matrix's data
 * pointer allocated.
 *
 * @param in the source matrix
 * @param out the destination matrix
 */

void copy(Matrix in, Matrix out) {
  dim3 block_size(16, 16);
  dim3 grid_size(INT_DIV_CEIL(in.sub.length, block_size.x),
                 INT_DIV_CEIL(in.length, block_size.y));
  copy_k<<<grid_size, block_size>>>(in, out);
}

